# -*- coding: utf-8 -*-
"""
统一分类数据集：ClassificationDataset
====================================

作者：zym1105（实现 & 中文注释）
日期：2025-08-19

功能概述
--------
本数据集类用于加载 retina / fundus 分类任务所需的数据。仅依赖一个入口：
- data_path: 目录下包含 annotations.json
- save_json_dir: 输出（标签映射、统计信息）保存目录

数据格式（annotations.json）
---------------------------
{
  "<image_name>.jpg": {
    "image_path": "/abs/or/rel/path/to/image.jpg",
    "image_label": "中文标签",
    ...  # 其他字段忽略
  },
  ...
}

要点
----
1. 中文标签会被统计频次，并按“出现次数从高到低”的顺序映射为整数标签：
   - 最多的类别 -> 0
   - 第二多 -> 1
   - 依此类推
2. 初始化时会落盘两类辅助信息：
   - diag_label_mapping.json : { "中文标签": int_label, ... }
   - label_stats.json : { "total": N, "per_class": { "中文标签": count, ... } }
3. __getitem__ 返回 (image_tensor, label_int, image_path)
4. 支持自定义 transform；未传入则使用常见 ImageNet 预处理

新增特性
--------
- 增加全局黑名单 `LABEL_BLACKLIST`（中文标签列表）。若样本的 `image_label` 在黑名单中，则直接过滤，不参与训练与统计。
  默认黑名单：
      ['不可用', '进一步检查', '其他眼底病']

zym1105 
25-08-19
"""

import os
import json
from collections import Counter
from typing import Dict, List, Optional, Tuple

from PIL import Image
from loguru import logger
import torch
from torch.utils.data import Dataset
from torchvision import transforms


# ===== 全局黑名单：这些中文标签对应的样本将被丢弃 =====
LABEL_BLACKLIST = {'不可用', '进一步检查', '其他眼底病'}


class ClassificationDataset(Dataset):
    """
    统一分类数据集：
    - 仅传入 data_path（内含 annotations.json）和 save_json_dir
    - 将中文标签映射为 int（按频次降序排序）
    - 落盘标签映射与统计信息
    - __getitem__ 返回 (image_tensor, label_int, image_path)

    参数
    ----
    data_path : str
        包含 annotations.json 的目录路径
    save_json_dir : str
        保存标签映射/统计文件的目录
    transform : Optional[torchvision.transforms.Compose]
        自定义图像预处理；若为 None，使用默认预处理

    属性
    ----
    data : List[Dict]
        每条记录包含 {"image_path": str, "label_text": str, "label": int}
    label2diag : Dict[int, str]
        整数标签 -> 中文标签
    diag2label : Dict[str, int]
        中文标签 -> 整数标签
    image_paths : List[str]
        方便外部访问的图片路径列表
    labels : List[int]
        方便外部访问的整数标签列表
    """

    def __init__(
        self,
        data_path: str,
        save_json_dir: str,
        transform: Optional[transforms.Compose] = None,
    ) -> None:
        os.makedirs(save_json_dir, exist_ok=True)

        # 1) 读取 annotations.json
        ann_file = os.path.join(data_path, "annotations.json")
        if not os.path.isfile(ann_file):
            raise FileNotFoundError(f"未找到 annotations.json：{ann_file}")

        with open(ann_file, "r", encoding="utf-8") as f:
            raw = json.load(f)

        # 2) 解析、清洗（只保留包含 image_path 与 image_label 的条目；并过滤黑名单标签）
        records: List[Dict] = []
        missing_cnt = 0
        blacklist_cnt = 0

        for _, item in raw.items():
            img_path = item.get("image_path")
            label_text = item.get("image_label")
            if not img_path or not label_text:
                missing_cnt += 1
                continue

            # 过滤黑名单标签
            # 说明：有些数据源可能带空格或全角空白，统一 strip 以避免误判
            label_text_norm = str(label_text).strip()
            if label_text_norm in LABEL_BLACKLIST:
                blacklist_cnt += 1
                continue

            # 统一成绝对路径，避免后续找不到文件
            img_path = os.path.abspath(img_path)
            records.append({"image_path": img_path, "label_text": label_text_norm})

        if len(records) == 0:
            raise ValueError(
                "annotations.json 中没有有效的数据条目（缺少 image_path/image_label，或被黑名单全部过滤）。"
            )

        logger.info(f"有效样本数：{len(records)}")
        logger.info(f"无效样本（缺字段）数：{missing_cnt}")
        logger.info(f"黑名单样本（被过滤）数：{blacklist_cnt}")
        if blacklist_cnt > 0:
            logger.info(f"当前黑名单：{sorted(list(LABEL_BLACKLIST))}")

        # 3) 统计中文标签频次，并构建“频次降序”的映射
        freq = Counter([r["label_text"] for r in records])
        # 按频次降序、同频次按标签文本排序稳定输出
        sorted_labels = sorted(freq.items(), key=lambda x: (-x[1], x[0]))
        diag2label: Dict[str, int] = {name: idx for idx, (name, _) in enumerate(sorted_labels)}
        label2diag: Dict[int, str] = {v: k for k, v in diag2label.items()}

        # 4) 应用映射，形成最终数据
        for r in records:
            r["label"] = diag2label[r["label_text"]]

        self.data = records
        self.diag2label = diag2label
        self.label2diag = label2diag
        self.image_paths = [r["image_path"] for r in self.data]
        self.labels = [r["label"] for r in self.data]

        # 5) 保存映射与统计
        mapping_path = os.path.join(save_json_dir, "diag_label_mapping.json")
        stats_path = os.path.join(save_json_dir, "label_stats.json")

        with open(mapping_path, "w", encoding="utf-8") as f:
            json.dump(self.diag2label, f, ensure_ascii=False, indent=2)

        stats_obj = {
            "total": len(self.data),
            "per_class": dict(freq),  # {中文标签: 计数}
            "blacklist": sorted(list(LABEL_BLACKLIST)),
            "filtered_by_blacklist": blacklist_cnt,
            "missing_or_invalid": missing_cnt,
        }
        with open(stats_path, "w", encoding="utf-8") as f:
            json.dump(stats_obj, f, ensure_ascii=False, indent=2)

        logger.success(f"标签映射已保存：{mapping_path}")
        logger.success(f"数据统计已保存：{stats_path}")

        # 6) 预处理 transform
        if transform is None:
            self.transform = transforms.Compose([
                transforms.Resize((224, 224)),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406],
                                     [0.229, 0.224, 0.225]),
            ])
        else:
            self.transform = transform

        # 7) 额外日志：展示前若干类的映射
        preview_k = min(10, len(self.diag2label))
        logger.info("前若干类映射预览（按频次降序）：")
        # 注意：diag2label 是 {中文: int}；这里为了按频次展示，用 sorted_labels
        for idx, (name, count) in enumerate(sorted_labels[:preview_k]):
            logger.info(f"  '{name}' -> {self.diag2label[name]}（样本数: {count}）")

    def __len__(self) -> int:
        return len(self.data)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int, str]:
        rec = self.data[idx]
        img_path = rec["image_path"]
        label_int = rec["label"]

        # 尽早捕获文件缺失问题
        if not os.path.isfile(img_path):
            raise FileNotFoundError(f"图像不存在：{img_path}")

        img = Image.open(img_path).convert("RGB")
        img_tensor = self.transform(img)
        return img_tensor, label_int, img_path

    # 可选：暴露一些便捷接口
    def get_label_text(self, label_int: int) -> str:
        """给定整数标签，返回中文标签"""
        return self.label2diag[label_int]

    def get_label_int(self, label_text: str) -> int:
        """给定中文标签，返回整数标签"""
        return self.diag2label[label_text]
